import re
from openai import OpenAI
import utils

"""
reflection的核心在于：对自己的产出结果进行检查，并结合实际输出情况进行优化和改进。
"""

api_key = '3dddcf7b2b836113433294985212a4f0.bvmegyRgPZcO7Pvj'
base_url = 'https://open.bigmodel.cn/api/paas/v4/'
model = 'glm-4-flash-250414'
client = OpenAI(api_key=api_key, base_url=base_url)


def generate_python_code(question: str, model: str = 'glm-4-flash-250414'):
    prompt = f'''你是一名python高级工程师，擅长通过编写python脚本来实现用户提出的需求，并且你编写的代码总是以最简洁、高效、注释明确的方式呈现。

    用户需求：{question}
    
    仅给出编写后的python代码，不能有其他任何内容。
'''

    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model=model,
        temperature=0,
    )
    content = response.choices[0].message.content.strip()
    return re.search(r"```python\s*([\s\S]*?)\s*```", content).group(1)


def refine_python_code(question: str, python_code: str, model: str = 'glm-4-flash-250414'):
    """
    Reflect on whether the outputs shown by the Python code answer the question.
    and propose improved python codes if needed.
    Returns (feedback, refined_python_code).
    """
    prompt = f'''你是一名python代码reviewer和refiner。
    
    用户的需求：{question}
    
    原始的pyhon代码：{python_code}
    
    第 1 步：简要评估该 Python 代码的输出是否完全回答了用户的问题。
    第 2 步：如果需要改进，请提供改进后的 Python 代码。
    如果原始 Python 代码已经正确，则保持不变返回。

    返回一个严格的 JSON，包含以下两个字段：
    {{
        "feedback": "<1-3 句话，用于解释存在的不足之处或确认代码的正确性。",
        "refined_python_code": "<最终版本的python代码>"
    }}

'''

    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )

    content = response.choices[0].message.content
    return utils.parse_content(content=content, python_code=python_code)


def refine_python_code_external_feedback(question: str, python_code: str, exec_output,
                                         model: str = 'glm-4-flash-250414'):
    """结合python执行的结果来优化。"""
    prompt = f'''你是一名python代码reviewer和refiner。

        用户的需求：{question}

        原始的pyhon代码：{python_code}
        python代码的执行结果：{exec_output}

        第 1 步：简要评估该 Python 代码的输出是否完全回答了用户的问题。
        第 2 步：如果需要改进（性能、可维护性、可读性、安全性、资源利用率等角度去考虑），请提供改进后的 Python 代码。
        如果原始 Python 代码无需改进优化，则保持不变返回。

        返回一个严格的 JSON，包含以下两个字段：
        {{
            "feedback": "<1-3 句话，用于解释存在的不足之处或确认代码的正确性。",
            "refined_python_code": "<最终版本的python代码>"
        }}

    '''

    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )

    content = response.choices[0].message.content
    return utils.parse_content(content=content, python_code=python_code)


def run_python_code_workflow(question: str, model: str = 'glm-4-flash-250414'):
    python_code = generate_python_code(question=question, model=model)
    exec_output = utils.exec_python_code(code=python_code)
    print(f'给出的第一版代码：{python_code}')
    print(f'\n第一版代码的执行结果：{exec_output["output"]}')

    print('*' * 100)

    feedback, refined_python_code = refine_python_code_external_feedback(question=question, python_code=python_code,
                                                                         exec_output=exec_output)
    exec_output = utils.exec_python_code(code=refined_python_code)
    print(f'给出的第二版代码：\n{refined_python_code}')
    print(f'\n第二版代码的执行结果：{exec_output["output"]}')

    return feedback, refined_python_code


if __name__ == '__main__':
    while True:
        question = input('请输入你的问题：')
        if question:
            feedback, refined_python_code = run_python_code_workflow(question=question)
            print(feedback)
            print(refined_python_code)

            print(f'最后的输出结果：{utils.exec_python_code(refined_python_code)}')
